from typing import Any, Iterable, Iterator, TypeVar, TypedDict

from zkl_aiutils_datasets.basics import Dataset, DatasetIterator, pause, resume
from zkl_aiutils_datasets.wrapping import wrap_dataset, wrap_dataset_iterator

AnyItem = TypeVar('AnyItem')
NoPadding = object()


class ChannelsBatchedIteratorState(TypedDict):
    iterator_state: tuple[Any, ...]
    channels_state: tuple[Any, ...]
    channels_item_state: tuple[Any, ...]


class ChannelsBatchedIterator(DatasetIterator[tuple[AnyItem, ...]]):
    def __init__(self,
        iterator: Iterator[Iterable[AnyItem]],
        channels_iterator: Iterable[Iterator[AnyItem]] | None = None,
        channels_state: Iterable[Any] | None = None,
        channels_n: int | None = None,
        pad_item: AnyItem = NoPadding,
    ):
        if channels_iterator is None:
            channels_iterator = []
            if channels_state is None:
                channels_state = []
            else:
                raise ValueError(f"Unexpected argument channels_state!")
            if channels_n is None:
                raise ValueError(f"Cannot resolve channels_n!")
        else:
            channels_iterator = list(channels_iterator)
            if channels_state is None:
                channels_state = None
            else:
                channels_state = list(channels_state)
                if len(channels_state) != len(channels_iterator):
                    raise ValueError(f"Unequal number of channels_state and channels_iterator!")
            if channels_n is None:
                channels_n = len(channels_iterator)

        self._iterator = iterator
        self._channels_iterator: list[Iterator[AnyItem] | None] = channels_iterator
        self._channels_state: list[Any] | None = channels_state
        self._channels_n = channels_n
        self._pad_item = pad_item

    def _reload_channel(self, channel_i: int | None):
        channel_state = None
        if self._channels_state is not None:
            try:
                channel_state = pause(self._iterator)
            except (TypeError, NotImplementedError):
                self._channels_state = None  # mark channels state unavailable

        try:
            channel_iterator = next(self._iterator)
            channel_iterator = wrap_dataset_iterator(channel_iterator)
        except StopIteration:
            channel_iterator = None
            if self._pad_item is NoPadding:
                raise

        if channel_i is None:
            self._channels_iterator.append(channel_iterator)
        else:
            self._channels_iterator[channel_i] = channel_iterator

        if self._channels_state is not None:
            if channel_i is None:
                self._channels_state.append(channel_state)
            else:
                self._channels_state[channel_i] = channel_state

        return channel_iterator

    def __next__(self) -> tuple[AnyItem, ...]:
        while len(self._channels_iterator) < self._channels_n:
            self._reload_channel(None)

        if all((iterator is None) for iterator in self._channels_iterator):
            raise StopIteration

        channels_item = []
        for channel_i in range(self._channels_n):
            while True:
                try:
                    channel_iterator = self._channels_iterator[channel_i]
                    channel_item = next(channel_iterator) \
                        if channel_iterator is not None else self._pad_item
                    channels_item.append(channel_item)
                    break
                except StopIteration:
                    if len(self._channels_iterator) > self._channels_n:
                        self._channels_iterator[channel_i] = self._channels_iterator.pop()
                        if self._channels_state is not None:
                            self._channels_state[channel_i] = self._channels_state.pop()
                        if all((iterator is None) for iterator in self._channels_iterator):
                            raise StopIteration
                    else:
                        new_channel_iterator = self._reload_channel(channel_i)
                        if new_channel_iterator is None:
                            if all((iterator is None) for iterator in self._channels_iterator):
                                raise StopIteration
        return tuple(channels_item)

    def __pause__(self) -> ChannelsBatchedIteratorState:
        if self._channels_state is None:
            raise TypeError(f"{self!r} does not support pausing!")
        channels_item_state = []
        for channel_iterator in self._channels_iterator:
            channel_item_state = pause(channel_iterator) if channel_iterator is not None else None
            channels_item_state.append(channel_item_state)
        return ChannelsBatchedIteratorState(
            iterator_state=pause(self._iterator),
            channels_state=tuple(self._channels_state),
            channels_item_state=tuple(channels_item_state))


class ChannelsBatchedDataset(Dataset[tuple[AnyItem, ...]]):
    def __init__(self,
        sequences: Iterable[Iterable[AnyItem]],
        channels_n: int,
        pad_item: AnyItem = NoPadding,
    ):
        self._sequences = sequences
        self._channels_n = channels_n
        self._pad_item = pad_item

    @property
    def channels_n(self) -> int:
        return self._channels_n

    @property
    def sequences(self) -> Iterable[Iterable[AnyItem]]:
        return self._sequences

    def __iter__(self) -> ChannelsBatchedIterator[AnyItem]:
        return ChannelsBatchedIterator(
            iterator=wrap_dataset_iterator(self._sequences),
            channels_n=self.channels_n,
            pad_item=self._pad_item)

    def __resume__(self, state: ChannelsBatchedIteratorState | None = None) -> ChannelsBatchedIterator[AnyItem]:
        channels_state = state['channels_state']
        channels_item_state = state['channels_item_state']

        channels_iterator = []
        sequences = wrap_dataset(self._sequences)
        for channel_state, channel_item_state in zip(channels_state, channels_item_state):
            try:
                channel_sequence = next(resume(sequences, channel_state))
                channel_sequence = wrap_dataset(channel_sequence)
                channel_iterator = resume(channel_sequence, channel_item_state)
            except StopIteration:
                channel_iterator = None
            channels_iterator.append(channel_iterator)

        return ChannelsBatchedIterator(
            iterator=resume(sequences, state['iterator_state']),
            channels_iterator=tuple(channels_iterator),
            channels_state=channels_item_state,
            channels_n=self.channels_n,
            pad_item=self._pad_item)
